import _thread as thread
import base64
import datetime
import hashlib
import hmac
import json
from urllib.parse import urlparse
import ssl
from datetime import datetime
from time import mktime
from urllib.parse import urlencode
from wsgiref.handlers import format_date_time
import websocket  # 使用websocket_client

# 全局变量用于存储回答
answer = ""

# 收到websocket消息的处理
def on_message(ws, message):
    data = json.loads(message)
    code = data['header']['code']
    if code != 0:
        ws.close()
    else:
        choices = data["payload"]["choices"]
        status = choices["status"]
        content = choices["text"][0]["content"]
        global answer
        answer += content
        if status == 2:
            ws.close()


# 收到websocket错误的处理
def on_error(ws, error):
    pass  # 移除调试输出

# 收到websocket关闭的处理
def on_close(ws, one, two):
    pass  # 移除调试输出

# 收到websocket连接建立的处理
def on_open(ws):
    thread.start_new_thread(run, (ws,))

def run(ws, *args):
    # 使用gen_params生成请求参数
    data = json.dumps(gen_params(appid=ws.appid, domain=ws.domain, message_history=ws.question))
    ws.send(data)

# 生成URL参数并返回
class Ws_Param(object):
    # 初始化
    def __init__(self, APPID, APIKey, APISecret, Spark_url):
        self.APPID = APPID
        self.APIKey = APIKey
        self.APISecret = APISecret
        self.host = urlparse(Spark_url).netloc
        self.path = urlparse(Spark_url).path
        self.Spark_url = Spark_url

    # 生成url
    def create_url(self):
        now = datetime.now()
        date = format_date_time(mktime(now.timetuple()))
        signature_origin = "host: " + self.host + "\n"
        signature_origin += "date: " + date + "\n"
        signature_origin += "GET " + self.path + " HTTP/1.1"

        signature_sha = hmac.new(self.APISecret.encode('utf-8'), signature_origin.encode('utf-8'),
                                 digestmod=hashlib.sha256).digest()
        signature_sha_base64 = base64.b64encode(signature_sha).decode(encoding='utf-8')
        authorization_origin = f'api_key="{self.APIKey}", algorithm="hmac-sha256", headers="host date request-line", signature="{signature_sha_base64}"'
        authorization = base64.b64encode(authorization_origin.encode('utf-8')).decode(encoding='utf-8')

        # 将请求的鉴权参数组合为字典
        v = {
            "authorization": authorization,
            "date": date,
            "host": self.host
        }
        # 拼接鉴权参数，生成url
        return self.Spark_url + '?' + urlencode(v)

# 生成请求参数
def gen_params(appid, domain, message_history):
    # 在原有的 message_history 中加入系统角色设定
    system_message = {
        "role": "system",
        "content": "你现在是外卖平台的客服，专门处理用户的订单问题，请根据用户的问题提供帮助。"
    }

    # 将系统设定的角色插入到 message_history 开头
    if isinstance(message_history, list):
        message_history.insert(0, system_message)

    data = {
        "header": {
            "app_id": appid,
            "uid": "1234"
        },
        "parameter": {
            "chat": {
                "domain": domain,
                "temperature": 0.5,
                "max_tokens": 1024,
                "auditing": "default"
            }
        },
        "payload": {
            "message": {
                "text": message_history  # 确保 message_history 是一个符合 API 格式的列表
            }
        }
    }
    return data


# 主函数，处理 websocket 连接
def main(appid, api_key, api_secret, Spark_url, domain, message_history):
    wsParam = Ws_Param(appid, api_key, api_secret, Spark_url)
    websocket.enableTrace(False)
    wsUrl = wsParam.create_url()
    ws = websocket.WebSocketApp(wsUrl, on_message=on_message, on_error=on_error, on_close=on_close, on_open=on_open)
    ws.appid = appid
    ws.question = message_history  # 使用原始的列表
    ws.domain = domain
    ws.run_forever(sslopt={"cert_reqs": ssl.CERT_NONE})

